# Shape and Topology optimization based on the phase field method (by: Akihiro Takezawa)
# PDE model: \int_D A*(\phi)e(u):e(v) dx = int_{\partial D} g\cdot v ds
#            ( A            if x \in Omega,
# A*(\phi) = { k(\phi)A     if x \in \xi,
#            ( k_min A      if x \in Omega_0
# AX = 2\mu X + \lamda tr(X)I 
# \lamda = (\nu) / ((1+v)*(1-2\nu)) E, \mu = 1 / (2(1+\nu)) E.
# min_{\phi} f(x) = int_{\partial D} g\cdot v ds
# L^2 gradient flow model: (\partial \phi)/(\partial t) = Kapa*\Delta \phi +
#            \phi*(1-\phi)*{\phi-1/2-30*\eta*(J'(\phi1))/(||J'(\phi1)||)*(1-\phi)*\phi)}

from ngsolve import *
from netgen.geom2d import SplineGeometry
from ngsolve.webgui import Draw
import numpy as np
## Setting mesh and Finite elements space

#LS=矩形半长边，SS=矩形半短边，FH=短边上分段的点
LS = 1
FH = 0.2
SS = 0.5
delta = 0.01
V  = 2*LS * 2*SS #total Volume
theta = 0.5      #using proportion
V0 = theta*V     #using volume
SIZE = 0.05      #mesh size                
                            #           y
# Set points of boundary    #           |
POINTS = [(LS, SS),         # 2---------|---------1
          (-LS,SS),         # |         |         |
          (-LS, FH),        # 3         |         8
          (-LS, -FH),       #-|---------|---------|---->x 
          (-LS, -SS),       # 4         |         7
          (LS, -SS),        # |         |         |
          (LS, -delta),     # 5-9-----11-12----10-6
          (LS, delta),
          (-LS+delta,-SS),
          (LS-delta,-SS),
          (-delta,-SS),
          (delta,-SS)]        
GEO = SplineGeometry()      

# Set name of points
P1,P2,P3,P4,P5,P6,P7,P8,P9,P10,P11,P12 = [GEO.AppendPoint(*P) for P in POINTS]
# Set name of boundaries
BNDS = [[["line",P1,P2],"BND1"], 
        [["line",P2,P3],"BND2"],
        [["line",P3,P4],"BND3"],
        [["line",P4,P5],"BND4"],
        [["line",P5,P9],"BND5"],
        [["line",P9,P11],"BND6"],
        [["line",P11,P12],"BND7"], # 牵引力作用于一小段
        [["line",P12,P10],"BND8"],
        [["line",P10,P6],"BND9"],
        [["line",P6,P7],"BND10"],
        [["line",P7,P8],"BND11"],
        [["line",P8,P1],"BND12"],]
# Add boundary to GEO
[GEO.Append(c,bc=bc) for c,bc in BNDS]
# Generate mesh
mesh = Mesh(GEO.GenerateMesh(maxh=SIZE))

# Set boundary condition
DB_MARK= "BND2|BND3|BND4"
NB_MARK= ""
forceBND = "BND11"

## Finite elements space of liner elasticity problem 
fesV = VectorH1(mesh, order = 1, dirichlet = DB_MARK)

## Finite elements space of Allen-Cahn equation 
fesS = H1(mesh, order = 1)

## Initial liner elasticity model parameter
nu      = 0.3
E       = 1
lamda   = nu / ((1+nu)*(1-2*nu))*E  #lame常数
mu      = 1 / (2*(1+nu))*E          #
p       = 3 #密度罚指数
phi0    = 0.0001 #密度下界
x0,y0 = LS,0 #外力作用点

# 设置CF类型的外力函数g
g = CF((0,-1)) #尽管设置在全区域，但只采用边界积分

# 设置体积力
force = CoefficientFunction( (0,0) )

# Set perimeter and volume penalize parameter
beta    = 0.01      #J^gamma = J + beta*F_gamma
alpha_k = 0.8       #F_gamma = gamma/2*\int_D |grad(phi)|^2 dx + 1/gamma\int_D W(phi) dx 
l_k     = 0.8         #L = J^gamma + lG + alpha_k/2 * G^2
gamma   = 0.01
xi      = 0.99      #alpha_{k+1} = alpha_k/xi

# control iterate
N = 20         # phase field update times
M = 3          # phase field compute times
tau = 0.035    #Compute time step

# beta 
cW = 0.1178  #cW = \int_phi0^1 sqrt(2W(s)) ds
#beta = beta / cW

## define absolution function
Absx = IfPos(x,x,-x)
Absy = IfPos(y,y,-y)

## Initial phi
phi = theta
phi_k = GridFunction(fesS)
phi_k.Set(phi)

## filter node and weight
rmin = 0.1  #filter radious
ft   = [1, 0]   #是否滤波：yes：1,-1, no: 0. -1表示反向加权的滤波 

if ft[0] == 1:
    # 初始化邻居关系表
    neighbors = {}
    # 初始化距离关系表
    neighbors_dist = {}
    neighbors_dist_sum = {}

    # 遍历所有点，找到每个点的邻居
    for vi in mesh.vertices:
        neighbors[vi.nr] = []  # 初始化每个成员的邻居列表
        neighbors_dist[vi.nr] = []
        neighbors_dist_sum[vi.nr] = rmin
        for vj in mesh.vertices:
            if vi.nr != vj.nr:  # 忽略自己
                dist = sqrt((vi.point[0]-vj.point[0])**2+(vi.point[1]-vj.point[1])**2)
                if dist < rmin:  # 距离小于 2
                    neighbors[vi.nr].append(vj.nr)
                    neighbors_dist[vi.nr].append(rmin - dist)
                    neighbors_dist_sum[vi.nr] += rmin - dist
elif ft[0] == -1:
    # 初始化邻居关系表
    neighbors = {}
    # 初始化距离关系表
    neighbors_dist = {}
    neighbors_dist_sum = {}

    # 遍历所有点，找到每个点的邻居
    for vi in mesh.vertices:
        neighbors[vi.nr] = []  # 初始化每个成员的邻居列表
        neighbors_dist[vi.nr] = []
        neighbors_dist_sum[vi.nr] = 0
        for vj in mesh.vertices:
            if vi.nr != vj.nr:  # 忽略自己
                dist = sqrt((vi.point[0]-vj.point[0])**2+(vi.point[1]-vj.point[1])**2)
                if dist < rmin:  # 距离小于 2
                    neighbors[vi.nr].append(vj.nr)
                    neighbors_dist[vi.nr].append(dist)
                    neighbors_dist_sum[vi.nr] += dist
    
'''
for key, value in neighbors.items():
    print(f"成员 {key} 的邻居是: {value}, 到邻居的距离之和: {neighbors_dist_sum[key]}")
'''     
## Strain and stress tensor
def strain(u):
    return Sym(Grad(u)) 
    # e(u) = 1/2(Grad(u)+Grad(u)^T)

def stress(u):
    return 2*mu*strain(u) + lamda*div(u)*Id(u.dim)
    # Aε = λ tr(ε)I + 2με , Ae(u) = 2μe(u) + λ div(u)I

## Volume constrain
def G(phi):
    return Integrate(phi,mesh) - V0

## Solve liner elasticity problem with phase field \phi
def solveElasticityProblem(phi):
    print("Being computing displacement field...")
    #define trial- and test- function
    u  = fesV.TrialFunction()   #trial function
    v = fesV.TestFunction()     #test function
    
    # the right hand side
    print("   Computing linear form...")
    f = LinearForm(phi*force*v*dx + g*v*ds(forceBND)).Assemble() 
        #\int_D f\cdot v dx + \int_Gamma_N g\cdot v ds
        
    # the stiffness matrix
    print("   Computing stiff matrix...")
    a = BilinearForm(fesV)
    a += phi**p*InnerProduct(stress(u), strain(v)) * dx
    a.Assemble()
        # int_D \rho^p Ae(u):e(v) dx
    
    gfu = GridFunction(fesV)  # 位移解
    gfu.vec.data = a.mat.Inverse(fesV.FreeDofs(),inverse="sparsecholesky")* f.vec
    
    return gfu

## Compute perimeter penalize
def F_gamma(phi):
    W = 0.25 * (phi-phi0)**2 * (phi-1)**2
    return gamma/2*Integrate(grad(phi)*grad(phi),mesh) + 1/gamma*Integrate(W,mesh)

## Compute objcetive functional
def ObjectFunctionVal(u,phi,l,alpha):
    print("Computing object function value...")
    L = Integrate(force*u,mesh) + Integrate(g*u,mesh.Boundaries(forceBND)) \
        + beta*F_gamma(phi) + l*G(phi) + alpha/2 * G(phi)**2
    return L

## Derivative of W
def dW(phi):
    ww = 0.5*(phi-phi0)*(phi-1)*(2*phi-1-phi0)
    return ww 

## Projection phase field function
def ProjectionPhaseField(phi, a, b):
    for i in range(fesS.ndof):
        phi.vec[i] = max(a,min(phi.vec[i],b))
    return phi

## Phase Field Filter
def PhaseFieldFilter(phi):
    for v in mesh.vertices:
        if ft[0] == 1:
            new = rmin * phi.vec[v.nr]
        elif ft[0] == -1:
            new = 0 
        for ner,ner_dist in zip(neighbors[v.nr], neighbors_dist[v.nr]):
            new += ner_dist*phi.vec[ner]
        phi.vec[v.nr] = new / neighbors_dist_sum[v.nr] 
    return phi 

## Update phase field function phi
def UpdatePhaseField(u,phi,l,alpha):
    print("Update phase field function...")
    deltaL = p*phi**(p-1)*InnerProduct(strain(u),stress(u)) \
            - beta/gamma*dW(phi) - (l + alpha*G(phi))
    #sensitivity filter
    if ft[1] == 1 or ft[1] == -1:
        gL = GridFunction(fesS)
        gL.Set(deltaL)
        deltaL = PhaseFieldFilter(gL)

    for i in range(M):
        rho, phy = fesS.TnT()
        f = LinearForm(fesS)
        f += 1/tau*phi*phy*dx + deltaL*phy*dx
        f.Assemble()
        a = BilinearForm(fesS)
        a += 1/tau*rho*phy*dx + beta*gamma*grad(rho)*grad(phy)*dx
        a.Assemble()
        phi.vec.data = a.mat.Inverse(fesS.FreeDofs()) * f.vec
        #Projection
        phi = ProjectionPhaseField(phi,phi0,1)
        #print(phi.vec)
    return phi

## Iter
value_L = []
value_C = []
value_F_gamma = []
value_V = []

for i in range(N):
    Draw(phi_k,mesh,"Phase Field")      #Draw optimization domain
    Fgamma_k    = beta*F_gamma(phi_k)   #Perimeter penalize
    Vol_k       = Integrate(phi_k,mesh) #Volume size
    u_k         = solveElasticityProblem(phi_k) #Solve state equation
    J_k         = ObjectFunctionVal(u_k,phi_k,l_k,alpha_k)  #Computing objective function val
    StrainE     = Integrate(phi_k**p*InnerProduct(stress(u_k),strain(u_k)),mesh)    #Computing strain energy
    l_k         = l_k + alpha_k*G(phi_k)    #Update the Lagrange multiplies
    print(f"{i+1:2} Iter: J = {J_k:.5}, \
        <Ae(u),e(u)> = {StrainE:.5},\
            F_gamma = {Fgamma_k:.5},\
                V =  {Vol_k:.5}.")
    value_L.append(J_k)
    value_C.append(StrainE)
    value_F_gamma.append(Fgamma_k)
    value_V.append(Vol_k)
    #Draw(u_k,mesh,"Displacement")   
    #Update Phase Field
    newphi_k       = UpdatePhaseField(u_k,phi_k,l_k,alpha_k)
    Draw(newphi_k,mesh)
    if ft[0] == 1 or ft[0] == -1:
        phi_k       = PhaseFieldFilter(newphi_k)
    else:
        phi_k       = newphi_k
    alpha_k     = alpha_k/xi

# 使用绝对路径（Windows 示例）
file_path = r"C:\Users\13285\Desktop\数值实验\PhaseField19.txt"

with open(file_path, 'w') as file:
    # 写入文件头部（可选）
    file.write(f"mesh: L={2*LS}, S={2*SS}, Dirichlet boundary={DB_MARK}, forcr BC={forceBND}\n")
    file.write(f"material: E={E}, nu={nu}\n")  
    file.write(f"parameter: \n  beta={beta},\n  alpha={alpha_k},\n  l={l_k},\n  gamma={gamma},\n \
        xi={xi},\n  N={N},\n  M={M},\n  tau={tau}\n")
    file.write('iter  L, C, F_gamma, V\n')
    
    # 使用 zip() 函数打包四个参数的数组
    for i,p1, p2, p3, p4 in zip(range(len(value_L)),value_L, value_C, value_F_gamma, value_V):
        # 将每组实验结果按行写入文件
        file.write(f"{i}, ||  {p1:.5},  ||  {p2:.5},  ||  {p3:.5},  ||  {p4:.5}\n")

print(f"实验结果已写入 {file_path} 文件")
